// 引入 js-pytorch 库中的 torch 对象，它包含了 PyTorch 的所有功能。
const { torch } = require("js-pytorch");

// 从 torch 中获取 nn 和 optim 模块，分别用于定义神经网络层和优化器。
const nn = torch.nn;
const optim = torch.optim;

// Define training hyperparameters:
// 定义训练的超参数。

const vocab_size = 52; // 词汇表的大小，即模型可以处理的不同词汇的总数。
const hidden_size = 32; // 隐藏层的大小，即隐藏层中神经元的数量。
const n_timesteps = 16; // 时间步长数，即序列数据中每个样本的长度。
const n_heads = 4; // 注意力机制中的头数。
const dropout_p = 0; // dropout概率，0 表示不使用dropout。
const batch_size = 8; // 每个训练批次的样本数量。

// Create Transformer decoder Module:
// 创建 Transformer 解码器模块。

class Transformer extends nn.Module {
  constructor(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p) {
    super(); // 调用父类 nn.Module 的构造函数。
    // Instantiate Transformer's Layers:
    // 实例化 Transformer 的各个层。

    this.embed = new nn.Embedding(vocab_size, hidden_size); // 词嵌入层。
    this.pos_embed = new nn.PositionalEmbedding(n_timesteps, hidden_size); // 位置编码层。
    this.b1 = new nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_p); // 第一个 Transformer 块。
    this.b2 = new nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_p); // 第二个 Transformer 块。
    this.ln = new nn.LayerNorm(hidden_size); // 层归一化。
    this.linear = new nn.Linear(hidden_size, vocab_size); // 最后的线性层，用于输出预测。
  }

  forward(x) {
    let z;
    // 前向传播函数，定义了数据如何通过网络层流动。

    z = torch.add(this.embed.forward(x), this.pos_embed.forward(x)); // 将词嵌入和位置编码相加。
    z = this.b1.forward(z); // 将结果传递给第一个 Transformer 块。
    z = this.b2.forward(z); // 将结果传递给第二个 Transformer 块。
    z = this.ln.forward(z); // 通过层归一化。
    z = this.linear.forward(z); // 通过最后的线性层得到预测结果。

    return z; // 返回模型的输出。
  }
}

// Instantiate your custom nn.Module:
// 实例化自定义的 nn.Module。

const model = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p); // 创建 Transformer 模型实例。

// Define loss function and optimizer:
// 定义损失函数和优化器。

const loss_func = new nn.CrossEntropyLoss(); // 交叉熵损失函数，常用于分类问题。
const optimizer = new optim.Adam(model.parameters(), (lr = 5e-3), (reg = 0)); // Adam 优化器，设置学习率和正则化参数。

// Instantiate sample input and output:
// 实例化样本输入和输出。

let x = torch.randint(0, vocab_size, [batch_size, n_timesteps, 1]); // 创建一个随机整数张量作为输入数据。
let y = torch.randint(0, vocab_size, [batch_size, n_timesteps]); // 创建一个随机整数张量作为目标数据。
let loss; // 用于存储损失值的变量。

// Training Loop:
// 训练循环。

for (let i = 0; i < 40; i++) {
  // Forward pass through the Transformer:
  // 通过 Transformer 进行前向传播。

  let z = model.forward(x); // 使用模型进行前向传播得到输出。

  // Get loss:
  // 计算损失。

  loss = loss_func.forward(z, y); // 使用损失函数计算模型输出和目标之间的损失。

  // Backpropagate the loss using torch.tensor's backward() method:
  // 使用 torch.tensor 的 backward() 方法进行反向传播。

  loss.backward(); // 反向传播损失，计算梯度。

  // Update the weights:
  // 更新权重。

  optimizer.step(); // 优化器根据计算出的梯度更新模型的权重。

  // Reset the gradients to zero after each training step:
  // 在每个训练步骤后将梯度重置为零。

  optimizer.zero_grad(); // 清除当前的梯度，为下一次梯度计算做准备。

  // Print loss at every iteration:
  // 在每次迭代时打印损失值。

  console.log(`Iter ${i} - Loss ${loss.data[0].toFixed(4)}`); // 打印当前迭代次数和损失值。
}